# simple_crystal.m 的体素实现!
# 或许体素更能体现原子的非刚体性和电子云（虽然精确的电子云分布是需要非常复杂的量子化学计算）
# Gitee Repo

import numpy as np
import plotly.graph_objects as go
from plotly.offline import plot

type = 2 # 晶体种类: 0:SC,1:BCC,2:FCC
Nx,Ny,Nz=3,3,3 # atom num 
Lx,Ly,Lz=10,10,10
dx = 0.2


# grid
x,y,z=np.meshgrid(np.arange(0,Lx,dx),np.arange(0,Ly,dx),np.arange(0,Lz,dx))
x,y,z=x.flatten(),y.flatten(),z.flatten()

# voxel data
u = np.zeros_like(x)
r = 1.0

def add_sphere(x,y,z,cx,cy,cz):
    cu = (x-cx)**2+(y-cy)**2+(z-cz)**2
    return np.exp(-cu)

X,Y,Z = np.meshgrid(np.arange(1,Nx+1),np.arange(1,Ny+1),np.arange(1,Nz+1))
XYZ = np.column_stack((X.flatten(),Z.flatten(),Y.flatten()))
del X,Y,Z

for ind in range(XYZ.shape[0]):
    i,j,k=XYZ[ind,0],XYZ[ind,1],XYZ[ind,2]
    if type == 0: #SC
        a = 2*r
        u += add_sphere(x,y,z,i*a,j*a,k*a);
        
    elif type == 1: #BCC
        a = 4.0/np.sqrt(3.0)*r
        u += add_sphere(x,y,z,i*a,j*a,k*a);
        
        if i != Nx and j != Ny and k != Nz: #edge
            u += add_sphere(x,y,z,(i+1/2)*a,(j+1/2)*a,(k+1/2)*a);
            
    elif type == 2: #FCC
        a = 4.0/np.sqrt(2.0)*r
        u += add_sphere(x,y,z,i*a,j*a,k*a);
        
        if not (j == Ny or k == Nz): #xedge
            u += add_sphere(x,y,z,i*a,(j+1/2)*a,(k+1/2)*a);
        
        if not (i == Nx or k == Nz): #yedge
            u += add_sphere(x,y,z,(i+1/2)*a,j*a,(k+1/2)*a);
        
        if not (i == Nx or j == Ny): #zedge
            u += add_sphere(x,y,z,(i+1/2)*a,(j+1/2)*a,k*a);
            
    else:
        raise NotImplementedError("This type is not implemented yet.")

            
fig = go.Figure(data=go.Volume(
    x=x,
    y=y,
    z=z,
    value=u,
    opacity=0.1, # needs to be small to see through all surfaces
    surface_count=20, # needs to be a large number for good volume rendering
    ))
plot(fig)

